import numpy as np
import soundfile as sf
from scipy.fft import dct, idct


def embed_watermark_in_channel(channel_data, watermark):
    """
    在单个声道中嵌入水印
    :param channel_data: 单声道音频数据
    :param watermark: 水印信息
    :return: 嵌入水印后的声道数据
    """
    watermark_bin = ''.join(format(ord(c), '08b') for c in watermark)
    watermark_len = len(watermark_bin)
    channel_len = len(channel_data)

    # 确保水印长度不超过音频容量
    if watermark_len * 32 > channel_len:
        print("水印过长，音频容量不足")
        return channel_data

    for i in range(watermark_len):
        start_index = i * 32
        end_index = start_index + 32
        if end_index > channel_len:
            break
        frame = channel_data[start_index:end_index]
        dct_frame = dct(frame)
        # 处理 DCT 系数为负值的情况
        rounded_dct_0 = int(np.round(dct_frame[0]))
        bit = int(watermark_bin[i])
        if rounded_dct_0 < 0:
            if (rounded_dct_0 & 1) != bit:
                dct_frame[0] = rounded_dct_0 + (1 if bit else -1)
        else:
            if (rounded_dct_0 & 1) != bit:
                dct_frame[0] = rounded_dct_0 + (1 if bit else -1)
        idct_frame = idct(dct_frame)
        channel_data[start_index:end_index] = idct_frame

    return channel_data


def embed_watermark(audio_path, watermark, output_path):
    """
    对音频文件按声道嵌入水印
    :param audio_path: 原始音频文件路径
    :param watermark: 水印信息
    :param output_path: 嵌入水印后音频文件的输出路径
    """
    try:
        audio, sample_rate = sf.read(audio_path)
    except Exception as e:
        print(f"读取音频文件时出错: {e}")
        return

    num_channels = audio.shape[1] if audio.ndim > 1 else 1
    if num_channels > 1:
        for channel in range(num_channels):
            audio[:, channel] = embed_watermark_in_channel(audio[:, channel], watermark)
    else:
        audio = embed_watermark_in_channel(audio, watermark)

    try:
        sf.write(output_path, audio, sample_rate)
        print(f"水印已成功嵌入到 {output_path}")
    except Exception as e:
        print(f"写入音频文件时出错: {e}")


def extract_watermark_from_channel(channel_data, watermark_length):
    """
    从单个声道中提取水印
    :param channel_data: 单声道音频数据
    :param watermark_length: 水印长度（比特数）
    :return: 提取的水印二进制字符串
    """
    watermark_bin = ""
    channel_len = len(channel_data)
    for i in range(watermark_length):
        start_index = i * 32
        end_index = start_index + 32
        if end_index > channel_len:
            break
        frame = channel_data[start_index:end_index]
        dct_frame = dct(frame)
        rounded_dct_0 = int(np.round(dct_frame[0]))
        bit = rounded_dct_0 & 1
        watermark_bin += str(bit)
    return watermark_bin


def extract_watermark(audio_path, watermark_length):
    """
    从音频文件中提取水印
    :param audio_path: 音频文件路径
    :param watermark_length: 水印长度（比特数）
    :return: 提取的水印信息
    """
    try:
        audio, sample_rate = sf.read(audio_path)
    except Exception as e:
        print(f"读取音频文件时出错: {e}")
        return ""

    num_channels = audio.shape[1] if audio.ndim > 1 else 1
    watermark_bin = ""
    if num_channels > 1:
        for channel in range(num_channels):
            channel_watermark_bin = extract_watermark_from_channel(audio[:, channel], watermark_length)
            watermark_bin += channel_watermark_bin
    else:
        watermark_bin = extract_watermark_from_channel(audio, watermark_length)

    # 将二进制水印转换为字符串
    watermark = ""
    for i in range(0, len(watermark_bin), 8):
        byte = watermark_bin[i:i + 8]
        if len(byte) == 8:
            try:
                watermark += chr(int(byte, 2))
            except ValueError:
                print(f"二进制转字符时出错，无效的二进制数据: {byte}")
    return watermark


# 示例使用
audio_path = "original.wav"
watermark = "SampleWatermark"
output_path = "output.wav"
embed_watermark(audio_path, watermark, output_path)

# 提取水印
extracted_watermark = extract_watermark(output_path, len(watermark) * 8)
print(f"提取的水印信息: {extracted_watermark}")
